#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Create by zhang
# Create on 2022/8/11 9:38
from enum import IntEnum
from typing import Tuple

from pandas import DataFrame
from numpy import nan

from core.dataClasses import StockTradeDataColumnName
from core.index_base import IndexBase


class CollectType(IntEnum):
    UP = 0
    DOWN = 1


class ChangeInValueIndex(IndexBase):
    """
    未来m-n（period[0], period[1]）天内涨跌幅度是否达到标准（dividing_boundary）
    """
    def __init__(self, adding_key:str, period:Tuple[int]=(1, 3), dividing_boundary:float=0.1,
                 collect_type:CollectType=CollectType.UP,
                 key=StockTradeDataColumnName.CLOSE,
                 remove_continued_label:bool=True):
        """

        :param adding_key: 标签列
        :param period: 收集标签的周期界限
        :param dividing_boundary: 收集标签分界涨跌幅度
        :param collect_type: 是上涨收集还是下跌收集
        :param key: 基准列
        @param remove_continued_label: 产生标签的时候是否删除连续出现的后续标签
        """
        super(ChangeInValueIndex, self).__init__()
        self.to_add_keys.append(adding_key)
        self.period:Tuple[int] = period
        self.dividing_boundary:float = dividing_boundary
        self.collect_type: CollectType = collect_type
        self.key:str = key
        self.remove_continued_label:bool = remove_continued_label

    def compute(self, data:DataFrame):
        self.added_keys.extend(self.to_add_keys)
        for col in self.to_add_keys:
            data[col] = nan
        is_continued = False
        for i in range(data.shape[0] - 1):
            # print(f"ChangeInValueIndex:{i}")
            # if i == 4066:
            #     print()
            target_index = i + self.period[1] if data.shape[0] - i - 1 > self.period[1] else data.shape[0]
            if i + 1 +  self.period[0] >= target_index:
                break
            if self.collect_type == CollectType.UP:
                max_idx = data.loc[i + 1 + self.period[0]:target_index, StockTradeDataColumnName.HIGH].idxmax()
                change = (data.loc[max_idx, StockTradeDataColumnName.HIGH] - data.loc[i, self.key]) / data.loc[i, self.key]
            elif self.collect_type == CollectType.DOWN:
                min_idx = data.loc[i + 1 + self.period[0]:target_index, StockTradeDataColumnName.LOW].idxmin()
                change = (data.loc[i, self.key] - data.loc[min_idx, StockTradeDataColumnName.HIGH]) / data.loc[i, self.key]
            label = 1 if change > self.dividing_boundary else 0
            if self.remove_continued_label:
                if label == 1 and is_continued:
                    data.loc[i, self.to_add_keys[0]] = 0
                elif label == 1 and not is_continued:
                    is_continued = True
                    data.loc[i, self.to_add_keys[0]] = 1
                else:
                    is_continued = False
                    data.loc[i, self.to_add_keys[0]] = 0
            else:
                data.loc[i, self.to_add_keys[0]] = label
        print(f"ChangeInValueIndex产生标签[{self.to_add_keys[0]}]数量:{data.loc[:, self.to_add_keys[0]].sum()}")